#!/usr/bin/env python3
"""
简化版教育LLM项目负载测试脚本
测试特定API的并发性能
"""

import asyncio
import aiohttp
import time
import random
from typing import List

# 测试配置
BASE_URL = "http://localhost:8000"
CONCURRENT_USERS = [100, 200, 300, 500]  # 不同并发级别
REQUESTS_PER_USER = 3  # 每个用户发送的请求数

# 测试数据
SAMPLE_QUESTIONS = [
    "什么是人工智能？",
    "机器学习的基本原理是什么？",
    "深度学习与传统机器学习有何区别？",
    "神经网络是如何工作的？",
    "自然语言处理技术有哪些应用？",
    "计算机视觉的主要任务是什么？",
    "强化学习的基本概念是什么？",
    "大数据处理的关键技术有哪些？",
    "云计算的服务模式有哪些？",
    "区块链技术的核心特点是什么？"
]

SAMPLE_SUBJECTS = ["数学", "物理", "化学", "生物", "历史", "地理", "语文", "英语", "政治", "计算机"]
SAMPLE_USERS = [f"user_{i}" for i in range(1, 101)]  # 100个测试用户

async def test_query_api(session, user_id: str, question: str) -> dict:
    """测试知识问答API"""
    url = f"{BASE_URL}/api/query"
    payload = {
        "question": question,
        "user_id": user_id
    }
    
    start_time = time.time()
    try:
        async with session.post(url, json=payload) as response:
            response_time = time.time() - start_time
            return {
                "success": response.status == 200,
                "status_code": response.status,
                "response_time": response_time
            }
    except Exception as e:
        response_time = time.time() - start_time
        return {
            "success": False,
            "status_code": 0,
            "response_time": response_time,
            "error": str(e)
        }

async def user_session(session, user_id: str, requests_count: int) -> List[dict]:
    """模拟单个用户会话"""
    results = []
    
    for _ in range(requests_count):
        # 随机选择一个问题
        question = random.choice(SAMPLE_QUESTIONS)
        
        # 测试问答API
        result = await test_query_api(session, user_id, question)
        results.append(result)
        
        # 随机小延迟
        await asyncio.sleep(random.uniform(0.1, 0.5))
    
    return results

async def run_load_test(concurrent_users: int) -> dict:
    """运行负载测试"""
    print(f"\n开始测试 {concurrent_users} 并发用户...")
    
    connector = aiohttp.TCPConnector(limit=concurrent_users, limit_per_host=concurrent_users//10)
    timeout = aiohttp.ClientTimeout(total=30)
    
    async with aiohttp.ClientSession(connector=connector, timeout=timeout) as session:
        start_time = time.time()
        
        # 创建所有用户任务
        tasks = []
        for i in range(concurrent_users):
            user_id = random.choice(SAMPLE_USERS)
            task = user_session(session, user_id, REQUESTS_PER_USER)
            tasks.append(task)
        
        # 并发执行所有任务
        all_user_results = await asyncio.gather(*tasks, return_exceptions=True)
        
        total_time = time.time() - start_time
    
    # 处理结果
    total_requests = 0
    successful_requests = 0
    failed_requests = 0
    response_times = []
    
    for user_results in all_user_results:
        if isinstance(user_results, Exception):
            print(f"用户会话出错: {user_results}")
            continue
            
        for result in user_results:
            total_requests += 1
            if result["success"]:
                successful_requests += 1
                response_times.append(result["response_time"])
            else:
                failed_requests += 1
    
    # 计算统计数据
    avg_response_time = sum(response_times) / len(response_times) if response_times else 0
    min_response_time = min(response_times) if response_times else 0
    max_response_time = max(response_times) if response_times else 0
    rps = total_requests / total_time if total_time > 0 else 0
    error_rate = failed_requests / total_requests if total_requests > 0 else 0
    
    return {
        "concurrent_users": concurrent_users,
        "total_requests": total_requests,
        "successful_requests": successful_requests,
        "failed_requests": failed_requests,
        "total_time": total_time,
        "avg_response_time": avg_response_time,
        "min_response_time": min_response_time,
        "max_response_time": max_response_time,
        "rps": rps,
        "error_rate": error_rate
    }

def print_test_result(result: dict):
    """打印测试结果"""
    print(f"\n{'='*50}")
    print(f"并发用户数: {result['concurrent_users']}")
    print(f"{'='*50}")
    print(f"总请求数: {result['total_requests']}")
    print(f"成功请求数: {result['successful_requests']}")
    print(f"失败请求数: {result['failed_requests']}")
    print(f"错误率: {result['error_rate']:.2%}")
    print(f"总耗时: {result['total_time']:.2f} 秒")
    print(f"平均响应时间: {result['avg_response_time']*1000:.2f} 毫秒")
    print(f"最小响应时间: {result['min_response_time']*1000:.2f} 毫秒")
    print(f"最大响应时间: {result['max_response_time']*1000:.2f} 毫秒")
    print(f"请求速率: {result['rps']:.2f} RPS")

async def main():
    """主函数"""
    print("教育LLM项目负载测试工具")
    print("基于 asyncio + aiohttp 实现")
    
    try:
        # 按不同并发级别进行测试
        for concurrent_users in CONCURRENT_USERS:
            result = await run_load_test(concurrent_users)
            print_test_result(result)
            
            # 在测试之间等待系统恢复
            if concurrent_users < CONCURRENT_USERS[-1]:
                print(f"\n等待10秒让系统恢复...")
                await asyncio.sleep(10)
                
    except KeyboardInterrupt:
        print("\n测试被用户中断")
    except Exception as e:
        print(f"测试过程中发生错误: {e}")

if __name__ == "__main__":
    asyncio.run(main())